#! /usr/bin/env python3
# -*-coding: utf-8-*-

"""
Created on Apr 20

@author: Moonkie

@attention: 
"""

import mxnet as mx
import numpy as np
import cv2
import os
import time


CURDIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.split(os.path.dirname(os.path.abspath(__file__)))[0]
DBDIR = os.path.join(CURDIR,'data')
ERROR_LOG_FILE = os.path.join(DBDIR,'error.log')


if not os.path.exists(DBDIR):
    os.makedirs(DBDIR)


class Bpnn(object):
    '''Bpnn neural network based on mxnet 

    '''
    def __init__(self,num_classes,num_feature,learn=.1,num_batch=10,ctx=mx.cpu()):
        data = mx.sym.Variable('data')
        net = mx.sym.FullyConnected(data,name='fc1',num_hidden=128)
        net = mx.sym.Dropout(net,name='dp1')
        net = mx.sym.Activation(data=net,name='relu1',act_type='sigmoid')
        net = mx.sym.FullyConnected(data=net,name='fc2',num_hidden=num_classes)
        net = mx.sym.SoftmaxOutput(data=net,name='out')
        self.net = net

        self.batch = num_batch
        self.classes = num_classes
        self.features = num_feature
        self.learn = learn

        self.exe = net.simple_bind(ctx=ctx,data=(self.batch,self.features))
        self.args = dict(zip(net.list_arguments(),self.exe.arg_arrays))
        for name,data in self.args.items():
            if 'weight' in name:
                data[:] = mx.random.uniform(.1,.03,data.shape)
        
    def forward(self,batch_input,is_train=False,ctx=mx.cpu()):
        self.exe.forward(is_train,data=batch_input)

    def propagation(self,out_label):
        self.args['out_label'] = out_label
        self.exe.backward()

        for weight,grad in zip(self.exe.arg_arrays,self.exe.grad_arrays):
            weight[:] -= self.learn * (grad/self.batch)
        # print(mx.nd.argmax_channel(self.exe.outputs[0]).asnumpy(),"error")
        # print(out_label,"outlabel")
        acc = (np.argmax(self.exe.outputs[0].asnumpy()) == out_label).sum()
        return acc

    def train(self,samples,labels,epoch=1000,is_train=True):
        # error = float(0)
        # for e in range(epoch):
        #     # batch_input = samples[:]
        #         # print(labels[i*self.batch:(i+1)*self.batch])
        #     # batch_label = labels[i*self.batch:(i+1)*self.batch]
        #     self.forward(batch_input=samples,is_train=is_train)
        #     acc = self.propagation(out_label=labels)
        #     print('[^_^] epoch %d acc is: %s' % (e,str(error/self.classes)))
        # print('[^_^] Finish trainning!')
        # print(self.args['out_label'],'\n','\n',self.exe.outputs[0].asnumpy())
        
        args = self.args
        ex = self.exe
        learning_rate = self.learn
        batch_size = self.batch
        for i in range(epoch):
            args['data'][:] = samples[:]
            ex.forward(is_train=True)
            args['out_label'][:] = labels[:]
            ex.backward()
            
            for weight,grad in zip(ex.arg_arrays,ex.grad_arrays):
                weight[:] -= learning_rate * (grad / batch_size)
            if i % 10 == 0:
                acc = (mx.nd.argmax_channel(ex.outputs[0]).asnumpy() == labels).sum()
                #print(mx.nd.argmax_channel(ex.outputs[0]).asnumpy())
                #exit()
                print("Current %d: %f" % (i,acc/labels.shape[0]))
        print(self.exe.outputs[0].asnumpy())

    def visualize(self):
        mx.viz.plot_network(self.net).view(Bpnn.__name__)
    
    @staticmethod
    def extract(lstfile,encode='utf8'):
        '''Extract lst file image and create output label

        '''
        data = []
        if not os.path.exists(lstfile):
            raise Exception("Nothing!! (extract/Bpnn.py)")
        label_file = np.load(lstfile)['label_file']
        for record in label_file:
            imgfilepath = record[1]
            img = cv2.imread(imgfilepath,cv2.IMREAD_UNCHANGED)
            if len(img.shape) == 3:
                img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
            img = Bpnn.binaryzation(img)
            img = cv2.resize(img,(92,112))
            img = Bpnn.dct(img)
            # show(img);exit(0)
            # img = img.flatten()
            # img = Bpnn.minmax(img)
            data.append(([np.argmax(record[0])],img))
        return np.array(data)
                

    @staticmethod
    def minmax(data):
        return (data - np.min(data))/(np.max(data) - np.min(data))
    
    @staticmethod
    def binaryzation(img):
        kernel = np.ones((2,2),np.uint8)
        if len(img.shape) == 3:
            img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        img = cv2.blur(img,(1,1))
        img = cv2.morphologyEx(img,cv2.MORPH_OPEN,kernel,iterations=2)
        img = cv2.threshold(img,70,255,cv2.THRESH_BINARY+cv2.THRESH_TRUNC)[1]
        thresh = np.mean(img)
        img[img>thresh] += 100
        img = cv2.equalizeHist(img)
        return 
    
    @staticmethod
    def dct(img):
        assert hasattr(img,"shape"),"Need a numpy array instance"
        assert len(img.shape) == 2,"img must be a gray image"
        img = cv2.medianBlur(img,3)
        data = np.array(img,np.float)/255   
        trans = cv2.dct(data)
        dct1 = np.uint8(trans)*255
        lwf = dct1[:12,:14]
        return lwf

    @staticmethod
    def split_label_data(data):
        assert len(data) != 0,"A empty array is given"
        assert  len(data[0]) == 2,"Need 2d array."
        labels = data[:,0]
        dat = data[:,1]
        print('[^_^] Split labels and data sunccessfully~')
        return np.array(list(labels),dtype=np.float32),np.array(list(dat))

    @staticmethod
    def getInitialParams(data,num_classes=6,num_batch=210):
        assert len(data.shape) == 2,"Need shape 2-d array."
        num_batch,num_feature =  data.shape
        return num_classes,num_feature,num_batch

def show(img):
    cv2.imshow('showing',img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    lst = os.path.join(CURDIR,'train.lst.npz')
    total_data = Bpnn.extract(lst)    
    labels,data = Bpnn.split_label_data(total_data)
    labels = labels.flatten()
    # print(labels)
    num_classes,num_feature,num_batch = Bpnn.getInitialParams(data)
    # print(data,labels)
    bp = Bpnn(num_classes=num_classes,num_feature=num_feature,num_batch=num_batch,ctx=mx.gpu())
    # bp.visualize()
    bp.train(data,labels)
    pass


